package map.core;

import map.core.util.GeoHelper;
import map.core.util.MercatorUtil;
import org.geotools.api.referencing.crs.CoordinateReferenceSystem;
import org.geotools.api.style.Style;
import org.geotools.coverage.grid.GridCoverage2D;
import org.geotools.geometry.jts.JTS;
import org.geotools.geometry.jts.ReferencedEnvelope;
import org.geotools.map.GridCoverageLayer;
import org.geotools.map.MapContent;
import org.geotools.referencing.CRS;
import org.geotools.renderer.GTRenderer;
import org.geotools.renderer.lite.StreamingRenderer;
import org.locationtech.jts.geom.Envelope;


import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.awt.image.WritableRaster;
import java.io.ByteArrayOutputStream;
import java.io.IOException;

public class MapRenderer {

   private static final CoordinateReferenceSystem CRS_4326;

    static {
        try {
            // CRS_4326 = CRS.parseWKT("GEOGCS[\"GCS_WGS_1984\",DATUM[\"D_WGS_1984\",SPHEROID[\"WGS_1984\",6378137.0,298.257223563]],PRIMEM[\"Greenwich\",0.0],UNIT[\"Degree\",0.0174532925199433]]");
            CRS_4326 = CRS.decode("epsg:4326");
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private   RasterLayer rasterLayer;

    private MapRenderer() {

    }


    public void addLayer(RasterLayer layer){
        this.rasterLayer = layer;
     }





    public byte[] getTile(int slice,RasterLayerStyle style,int level, int row, int column) {

        try {
            int imageWidth = 256;
            int imageHeight = 256;
            Envelope envelope = MercatorUtil.getGeographyBox(level, row, column, 256);
            ReferencedEnvelope mapBounds;
            mapBounds = new ReferencedEnvelope(envelope, CRS_4326);


            BufferedImage bufferedImage = getImage(slice,style, imageWidth, imageHeight, mapBounds);
            ByteArrayOutputStream rs = new ByteArrayOutputStream();
            ImageIO.write(bufferedImage, "png", rs);


            return rs.toByteArray();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }


    public byte[] getImage(int slice,RasterLayerStyle style,double scale){
        return  this.getImage(slice,style,scale,this.rasterLayer.getBounds());
    }

    public byte[] getImage(int slice,RasterLayerStyle style,double scale,LonLatBounds bounds){

        if(bounds == null){
            throw new RuntimeException("bounds 不可为空");
        }

        try {
            if(bounds == null){
                bounds = this.rasterLayer.getBounds();
            }

            LonLatBounds fullBounds = this.rasterLayer.getBounds();
            int fullImageWidth =  this.rasterLayer.getWidth();
            int fullImageHeight =  this.rasterLayer.getHeight();
            double widthResolution =  fullBounds.getLonSpan()/fullImageWidth ;
            //   double heightResolution =  fullImageHeight/fullBounds.getLatSpan() ;

            int imageWidth = (int) (bounds.getLonSpan() / widthResolution* scale);
            int imageHeight = (int) (bounds.getLatSpan() / widthResolution * scale);



            ReferencedEnvelope envelope = new ReferencedEnvelope(bounds.getMinLon(),bounds.getMaxLon(),bounds.getMinLat(),bounds.getMaxLat(),CRS_4326);


            BufferedImage bufferedImage = getImage(slice,style, imageWidth, imageHeight, envelope);
            ByteArrayOutputStream rs = new ByteArrayOutputStream();
            ImageIO.write(bufferedImage, "png", rs);
            return rs.toByteArray();

        } catch (IOException e) {
            throw new RuntimeException(e);
        }

    }

    private BufferedImage genBufferedImage(GridCoverage2D dataGrid,Style rasterStyle, int imageWith, int imageHeight, ReferencedEnvelope mapBounds){
        try{

            MapContent map = new MapContent();
            map.getViewport().setCoordinateReferenceSystem(GeoHelper.WEB_MERCATOR);
            GridCoverageLayer rasterLayer = new GridCoverageLayer( dataGrid,rasterStyle);
            map.addLayer(rasterLayer);
            GTRenderer renderer = new StreamingRenderer();
            renderer.setMapContent(map);
            Rectangle imageBounds = new Rectangle(0, 0, imageWith, imageHeight);
            BufferedImage image = new BufferedImage(imageWith, imageHeight, BufferedImage.TYPE_INT_ARGB);

            Graphics2D gr = image.createGraphics();
            renderer.paint(gr, imageBounds, JTS.transform(mapBounds, GeoHelper.WGS84_TO_WEB_MERCATOR));
            gr.dispose();

            map.dispose();
            return image;
        }
        catch (Exception e){
            throw new RuntimeException(e);
        }
    }



    private BufferedImage getImage(int slice,RasterLayerStyle style,int imageWith, int imageHeight, ReferencedEnvelope mapBounds) {
        GridCoverage2D data = this.rasterLayer.getGridData(slice);
         BufferedImage baseImage = genBufferedImage(data,style.getBaseStyle(),imageWith,imageHeight,mapBounds);
         BufferedImage noDataImage = genBufferedImage(data,style.getNoDataStyle(),imageWith,imageHeight,mapBounds);

        WritableRaster baseImageRaster = baseImage.getRaster();
        WritableRaster noDataImageRaster = noDataImage.getRaster();


        int[] noDataColor  = style.getNoDataRGBA();

        int[] baseImageValue = new int[4];
        int[] noDataImageValue = new int[4];

        for(int i =  0  ; i < imageWith ; i ++){
            for(int j = 0 ; j < imageHeight ; j ++){
                baseImageRaster.getPixel(i,j,baseImageValue);
                noDataImageRaster.getPixel(i,j,noDataImageValue);
                if(noDataImageValue[0] >0){//当前像素为NoData
                    baseImageRaster.setPixel(i,j,noDataColor);
                }
            }
        }

        return baseImage;
    }




    public static MapRenderer create(RasterLayer layer){
        MapRenderer renderer = new MapRenderer();
        renderer.addLayer(layer);
        return  renderer;
    }

 }
